#include "opencv_undistort.h"

namespace casia {

    OpencvUnDistort::OpencvUnDistort() = default;

    OpencvUnDistort::OpencvUnDistort(const std::string &parameterFilePath) {
        if (!fileExists(parameterFilePath)) {
            spdlog::error("OpencvUnDistort Object Instance failed, file [{}] doesn't exists.", parameterFilePath);
            return;
        }
        YAML::Node config = YAML::LoadFile(parameterFilePath);
        /// 矩阵正则
        std::regex regex(R"(-?\d+(\.\d+)?)", std::regex_constants::ECMAScript);
        ///  解析内参矩阵
        auto intrinsicMatrixYaml = config["intrinsic_matrix"].as<std::vector<std::string>>();
        for (int row = 0; row < intrinsicMatrixYaml.size(); row++) {
            spdlog::debug("Get intrinsic matrix row {}: {}", row, intrinsicMatrixYaml[row]);
            std::string item = intrinsicMatrixYaml[row];
            std::regex_search(item, regex);
            auto begin = std::sregex_iterator(item.begin(), item.end(), regex);
            auto end = std::sregex_iterator();
            int col = 0;
            for (std::sregex_iterator i = begin; i != end; ++i, col++) {
                std::smatch match = *i;
                std::string res = match.str();
                spdlog::debug("Get intrinsic parameter {}", res);
                this->intrinsicMatrix.at<double>(row, col) = (double) stod(res);
            }
        }
        /// 解析畸变矩阵
        auto distortionMatrixYaml = config["distortion_matrix"].as<std::string>();
        spdlog::debug("Get distortion matrix {}", distortionMatrixYaml);
        std::regex_search(distortionMatrixYaml, regex);
        auto begin = std::sregex_iterator(distortionMatrixYaml.begin(), distortionMatrixYaml.end(), regex);
        auto end = std::sregex_iterator();
        int col = 0;
        for (std::sregex_iterator i = begin; i != end; ++i, col++) {
            std::smatch match = *i;
            std::string res = match.str();
            spdlog::debug("Get distortion parameter {}", res);
            this->distortionMatrix.at<double>(0, col) = (double) stod(res);
        }
    }

    OpencvUnDistort::~OpencvUnDistort() {
        spdlog::debug("OpencvUnDistort Destructor Function call.");
    }

    bool OpencvUnDistort::loadImage(const std::string &path) {
        if (!fileExists(path)) {
            spdlog::error("OpencvUnDistort loadImage failed, file [{}] doesn't exists.", path);
            return false;
        }
        this->image = cv::imread(path);
        if (this->image.cols > 0) {
            return true;
        } else {
            spdlog::error("OpencvUnDistort loadImage failed, file [{}] may not a valid image.", path);
            return false;
        }
    }

    void showImage(const cv::Mat &image) {
        if (image.cols > 0) {
            std::string title = std::string("Image ").append(std::to_string(image.cols)).append(" * ").append(
                    std::to_string(image.rows));
            cv::imshow(title, image);
            cv::waitKey(0);
        }
    }

    void OpencvUnDistort::showOriginImage() {
        showImage(this->image);
    }

    void OpencvUnDistort::showUnDistortImage() {
        showImage(this->unDistortImage);
    }

    void OpencvUnDistort::UnDistort() {
        /// 裁剪掉黑边
//        cv::undistort(this->image, this->unDistortImage, this->intrinsicMatrix, this->distortionMatrix);
        /// 不裁剪黑边
        cv::Mat map1, map2;
        cv::Size imageSize;
        imageSize = this->image.size();
        initUndistortRectifyMap(this->intrinsicMatrix, this->distortionMatrix, cv::Mat(),
                                getOptimalNewCameraMatrix(this->intrinsicMatrix, this->distortionMatrix, imageSize, 1,
                                                          imageSize, nullptr),
                                imageSize, CV_16SC2, map1, map2);
        remap(this->image, this->unDistortImage, map1, map2, cv::INTER_LINEAR);
    }

    void OpencvUnDistort::saveUnDistort(const std::string &path) {
        if (this->unDistortImage.cols > 0) {
            cv::imwrite(path, this->unDistortImage);
        }
    }

}

int main(int argc, char** argv) {
    spdlog::set_level(spdlog::level::debug);

    if (argc != 3) {
        spdlog::error("Incorrect, to use this program, the command should like this: ./undistort <inputImagePath> <outputImagePath>");
        return 0;
    }
    std::string path = R"(/mnt/e/workspace-clion/undistort/resources/camera_parameters.yaml)";
    casia::OpencvUnDistort unDistort(path);
    unDistort.loadImage(std::string(argv[1]));
//    unDistort.showOriginImage();
    unDistort.UnDistort();
    unDistort.saveUnDistort(std::string(argv[2]));
    return 0;
}